import torch                         # 导入 PyTorch 主包
import torch.nn as nn                # 导入神经网络模块
import torch.nn.functional as F      # 导入函数式 API（激活、loss 等）
from torchvision import models       # 导入 torchvision 的模型（VGG 等）
from torchvision.ops import RoIPool, nms  # 导入 RoI Pooling 和 NMS 操作


# ========= 一些 bbox（边界框）工具函数 =========

def box_area(boxes):
    """
    计算一组 boxes 的面积
    boxes: (N, 4)，格式为 (x1, y1, x2, y2)
    """
    w = (boxes[:, 2] - boxes[:, 0]).clamp(min=0)  # 右减左得到宽度，并限制最小值为 0，防止负数
    h = (boxes[:, 3] - boxes[:, 1]).clamp(min=0)  # 下减上得到高度，并限制最小值为 0
    return w * h                                   # 面积 = 宽 * 高


def box_iou(boxes1, boxes2):
    """
    计算两组 boxes 的 IoU
    boxes1: (N, 4)
    boxes2: (M, 4)
    返回: (N, M) 的 IoU 矩阵
    """
    area1 = box_area(boxes1)                      # 计算第一组 boxes 的面积 (N,)
    area2 = box_area(boxes2)                      # 计算第二组 boxes 的面积 (M,)

    # 交集左上角坐标：取两者左上角坐标的最大值
    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])   # (N, M, 2)
    # 交集右下角坐标：取两者右下角坐标的最小值
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])   # (N, M, 2)
    # 交集宽高：右下 - 左上，并限制最小为 0
    wh = (rb - lt).clamp(min=0)                          # (N, M, 2)
    # 交集面积 = 宽 * 高
    inter = wh[:, :, 0] * wh[:, :, 1]                    # (N, M)

    # 并集面积 = area1 + area2 - 交集
    union = area1[:, None] + area2 - inter
    # IoU = 交集 / 并集（避免除 0）
    iou = inter / union.clamp(min=1e-6)
    return iou


def encode_boxes(proposals, gt_boxes):
    """
    生成回归目标（t_x, t_y, t_w, t_h）
    proposals: (N, 4)，原始候选框
    gt_boxes:  (N, 4)，与 proposals 一一对应的 GT 框
    """
    # proposals 的中心坐标和宽高
    px = (proposals[:, 0] + proposals[:, 2]) / 2.0       # 计算 proposal 中心 x
    py = (proposals[:, 1] + proposals[:, 3]) / 2.0       # 计算 proposal 中心 y
    pw = (proposals[:, 2] - proposals[:, 0]).clamp(min=1e-6)  # 计算 proposal 宽度并避免除 0
    ph = (proposals[:, 3] - proposals[:, 1]).clamp(min=1e-6)  # 计算 proposal 高度并避免除 0

    # gt_boxes 的中心坐标和宽高
    gx = (gt_boxes[:, 0] + gt_boxes[:, 2]) / 2.0         # 计算 GT 中心 x
    gy = (gt_boxes[:, 1] + gt_boxes[:, 3]) / 2.0         # 计算 GT 中心 y
    gw = (gt_boxes[:, 2] - gt_boxes[:, 0]).clamp(min=1e-6)    # 计算 GT 宽度
    gh = (gt_boxes[:, 3] - gt_boxes[:, 1]).clamp(min=1e-6)    # 计算 GT 高度

    # 位置偏移：中心点差值除以 proposal 的宽高
    tx = (gx - px) / pw                                  # x 方向偏移
    ty = (gy - py) / ph                                  # y 方向偏移
    # 尺度偏移：宽高比取 log
    tw = torch.log(gw / pw)                              # 宽度缩放
    th = torch.log(gh / ph)                              # 高度缩放

    # 堆叠成 (N, 4) 的回归目标向量
    return torch.stack([tx, ty, tw, th], dim=1)


def decode_boxes(proposals, deltas):
    """
    根据回归结果（deltas）从 proposals 解码出预测 bbox
    proposals: (N, 4)
    deltas:    (N, 4)  [t_x, t_y, t_w, t_h]
    """
    # proposals 的中心和宽高
    px = (proposals[:, 0] + proposals[:, 2]) / 2.0       # proposal 中心 x
    py = (proposals[:, 1] + proposals[:, 3]) / 2.0       # proposal 中心 y
    pw = (proposals[:, 2] - proposals[:, 0]).clamp(min=1e-6)  # proposal 宽度
    ph = (proposals[:, 3] - proposals[:, 1]).clamp(min=1e-6)  # proposal 高度

    # 回归偏移量
    tx = deltas[:, 0]                                   # x 偏移量
    ty = deltas[:, 1]                                   # y 偏移量
    tw = deltas[:, 2]                                   # 宽度缩放量
    th = deltas[:, 3]                                   # 高度缩放量

    # 解码出预测框中心和宽高
    gx = tx * pw + px                                   # 预测中心 x
    gy = ty * ph + py                                   # 预测中心 y
    gw = torch.exp(tw) * pw                             # 预测宽度
    gh = torch.exp(th) * ph                             # 预测高度

    # 转回 (x1, y1, x2, y2) 格式
    x1 = gx - gw / 2.0                                  # 左上 x
    y1 = gy - gh / 2.0                                  # 左上 y
    x2 = gx + gw / 2.0                                  # 右下 x
    y2 = gy + gh / 2.0                                  # 右下 y

    return torch.stack([x1, y1, x2, y2], dim=1)         # 返回 (N, 4)


def clip_boxes_to_image(boxes, height, width):
    """
    将 boxes 限制在图像边界内
    boxes: (N, 4)
    height, width: 图像高和宽
    """
    boxes[:, 0] = boxes[:, 0].clamp(min=0, max=width - 1)   # 限制 x1 在 [0, w-1]
    boxes[:, 1] = boxes[:, 1].clamp(min=0, max=height - 1)  # 限制 y1 在 [0, h-1]
    boxes[:, 2] = boxes[:, 2].clamp(min=0, max=width - 1)   # 限制 x2 在 [0, w-1]
    boxes[:, 3] = boxes[:, 3].clamp(min=0, max=height - 1)  # 限制 y2 在 [0, h-1]
    return boxes                                            # 返回裁剪后的 boxes


# ========= GT 匹配 & 目标生成 =========

def assign_targets_to_proposals(proposals, gt_boxes, gt_labels,
                                num_classes,
                                pos_iou_thresh=0.5,
                                neg_iou_thresh=0.1):
    """
    对一个图像内的所有 proposals 做 GT 匹配，生成训练用目标：
      - labels: (N,)  0=背景，1..K=类别
      - bbox_targets: (N, 4) 回归目标（仅对正样本有效）
    proposals: (N, 4)
    gt_boxes:  (M, 4)
    gt_labels: (M,) 值在 1..K
    num_classes: 类别数（含背景，通常用不到这里）
    pos_iou_thresh: IoU >= 该阈值视为正样本
    neg_iou_thresh: IoU < 该阈值视为忽略（负样本下界）
    """
    device = proposals.device                            # 获取当前张量所在设备
    N = proposals.shape[0]                               # proposals 数量

    labels = torch.zeros((N,), dtype=torch.long, device=device)       # 初始化标签为 0（背景）
    bbox_targets = torch.zeros((N, 4), dtype=torch.float32, device=device)  # 初始化回归目标为 0

    if gt_boxes.numel() == 0:                            # 如果没有 GT
        # 没有 GT 时，全部 proposal 都视为背景
        return labels, bbox_targets

    # 计算每个 proposal 和每个 GT 的 IoU，得到 (N, M) 矩阵
    ious = box_iou(proposals, gt_boxes)                  # IoU 矩阵
    max_iou, max_idx = ious.max(dim=1)                   # 对每个 proposal 找到最大 IoU 及对应 GT 索引

    # 正样本：IoU >= 正阈值
    pos_mask = max_iou >= pos_iou_thresh
    # 负样本：neg_iou_thresh <= IoU < pos_iou_thresh（中间区间，具体策略视实现而定）
    neg_mask = (max_iou < pos_iou_thresh) & (max_iou >= neg_iou_thresh)

    # 给负样本标签设为 0（背景）
    labels[neg_mask] = 0
    # 给正样本标签设成对应 GT 的类别（1..K）
    labels[pos_mask] = gt_labels[max_idx[pos_mask]]

    # 对正样本生成 bbox 回归目标
    if pos_mask.any():                                   # 如果存在正样本
        matched_gt_boxes = gt_boxes[max_idx[pos_mask]]   # 取出与正样本匹配的 GT 框
        matched_props = proposals[pos_mask]              # 取出对应的 proposals
        # 计算回归目标（t_x, t_y, t_w, t_h）
        targets_pos = encode_boxes(matched_props, matched_gt_boxes)  # (num_pos, 4)
        bbox_targets[pos_mask] = targets_pos             # 填入对应位置

    return labels, bbox_targets                          # 返回标签和回归目标


# ========= VGG Backbone =========

class VGGBackbone(nn.Module):
    """
    使用 VGG16 的卷积部分作为特征提取 backbone
    """
    def __init__(self):
        super().__init__()                               # 调用父类构造函数
        vgg = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)  # 加载 ImageNet 预训练 VGG16
        self.features = vgg.features                     # 只保留卷积层部分（features）

    def forward(self, x):
        return self.features(x)                          # 输入图像 x，输出特征图 (N, C, H', W')


# ========= Fast R-CNN Head =========

class FastRCNNHead(nn.Module):
    """
    Fast R-CNN head：
      - RoI Pooling
      - 两层全连接
      - 分类 & 边框回归头
    """
    def __init__(self, in_channels, roi_output_size, num_classes):
        super().__init__()                               # 调用父类构造函数
        self.num_classes = num_classes                   # 类别数（含背景）
        self.roi_output_size = roi_output_size           # RoI Pooling 的输出空间尺寸，例如 (7, 7)

        # RoI Pooling 层：从特征图上根据 RoI 坐标截取固定大小的特征
        self.roi_pool = RoIPool(
            output_size=roi_output_size,                 # 输出空间大小（高, 宽）
            spatial_scale=1.0 / 16.0                     # 特征图相对于原图的缩放比例（VGG16 总 stride = 16）
        )

        # RoI Pooling 输出的特征维度：通道数 * 高 * 宽
        self.fc_input_dim = in_channels * roi_output_size[0] * roi_output_size[1]
        # 第一层全连接，将特征映射到 4096 维
        self.fc1 = nn.Linear(self.fc_input_dim, 4096)
        # 第二层全连接，继续保持 4096 维
        self.fc2 = nn.Linear(4096, 4096)

        self.relu = nn.ReLU(inplace=True)               # ReLU 激活函数
        self.dropout = nn.Dropout(p=0.5)                # Dropout 防止过拟合

        # 分类头：输出 num_classes 维（含背景）
        self.cls_score = nn.Linear(4096, num_classes)
        # 回归头：对每个类别输出 4 个 bbox 偏移量 → num_classes * 4
        self.bbox_pred = nn.Linear(4096, num_classes * 4)

        self._init_weights()                            # 初始化权重

    def _init_weights(self):
        """
        对全连接层进行简单的高斯初始化
        """
        for m in [self.fc1, self.fc2, self.cls_score, self.bbox_pred]:
            nn.init.normal_(m.weight, std=0.01)         # 权重初始化为均值 0、方差 0.01 的正态分布
            nn.init.constant_(m.bias, 0.)               # 偏置初始化为 0

    def forward(self, feature_map, rois):
        """
        feature_map: (N, C, H', W')，来自 backbone 的特征图
        rois: (num_rois, 5)，每行为 [batch_idx, x1, y1, x2, y2]（原图尺度坐标）
        """
        # 1. 在特征图上对所有 RoI 做 RoI Pooling，得到 (num_rois, C, 7, 7)
        pooled = self.roi_pool(feature_map, rois)
        # 2. 将每个 RoI 的特征展平为一维向量 (num_rois, C*7*7)
        x = pooled.flatten(start_dim=1)

        # 3. 全连接层 + ReLU + Dropout
        x = self.fc1(x)                                 # 通过第一层 FC
        x = self.relu(x)                                # ReLU 激活
        x = self.dropout(x)                             # Dropout

        x = self.fc2(x)                                 # 第二层 FC
        x = self.relu(x)                                # ReLU 激活
        x = self.dropout(x)                             # Dropout

        # 4. 得到分类 logits 和 bbox 回归输出
        cls_logits = self.cls_score(x)                  # 分类 logits，形状 (num_rois, num_classes)
        bbox_deltas = self.bbox_pred(x)                 # 回归输出，形状 (num_rois, num_classes * 4)

        return cls_logits, bbox_deltas                  # 返回分类与回归结果


# ========= 整体 Fast R-CNN 模型 =========

class FastRCNN(nn.Module):
    """
    整体 Fast R-CNN 模型，包含：
      - Backbone: VGG16 卷积层
      - Head: RoI Pooling + FC + 分类 & 回归
    支持：
      - 训练模式：传入 images, rois, targets，返回 loss 字典
      - 推理模式：传入 images, rois，返回检测结果（boxes, labels, scores）
    """
    def __init__(self, num_classes=21, roi_output_size=(7, 7)):
        super().__init__()                               # 调用父类构造
        self.num_classes = num_classes                   # 类别数（含背景）
        self.backbone = VGGBackbone()                    # VGG16 特征提取 backbone
        self.head = FastRCNNHead(                        # Fast R-CNN head
            in_channels=512,                             # VGG16 conv5 输出通道数 512
            roi_output_size=roi_output_size,             # RoI Pooling 输出尺寸 (7,7)
            num_classes=num_classes                      # 类别数
        )

        # 分类损失：多类交叉熵
        self.cls_loss_func = nn.CrossEntropyLoss()
        # 回归损失：Smooth L1（带 beta）
        self.reg_loss_func = nn.SmoothL1Loss(beta=1.0)

    def forward(self, images, rois, targets=None,
                score_thresh=0.05, nms_thresh=0.5, topk_per_image=100):
        """
        images: (N, 3, H, W) 的输入图像 batch
        rois: (num_rois, 5)，每行为 [batch_idx, x1, y1, x2, y2]
        targets:
            - 训练时传入: list[dict]，长度 N
                每个 dict 包含:
                    "boxes": (Mi, 4)
                    "labels": (Mi,)
            - 推理时传入: None
        score_thresh: 推理时的置信度阈值
        nms_thresh: NMS 的 IoU 阈值
        topk_per_image: 每张图最多保留的检测数
        """
        device = images.device                           # 获取当前图像所在设备
        feature_map = self.backbone(images)             # 通过 backbone 提取特征图

        cls_logits, bbox_deltas = self.head(feature_map, rois)  # 通过 head 得到分类和回归输出

        if targets is not None:
            # ====== 训练模式：计算 loss ========
            return self.compute_loss(
                rois, cls_logits, bbox_deltas, targets
            )
        else:
            # ====== 推理模式：NMS & 输出检测框 ========
            return self.inference(
                rois, cls_logits, bbox_deltas,
                images_shapes=[(images.shape[2], images.shape[3])] * images.shape[0],
                score_thresh=score_thresh,
                nms_thresh=nms_thresh,
                topk_per_image=topk_per_image
            )

    # --------- Loss 计算 ---------
    def compute_loss(self, rois, cls_logits, bbox_deltas, targets):
        """
        计算 Fast R-CNN 的分类和回归损失
        rois: (num_rois, 5)
        cls_logits: (num_rois, num_classes)
        bbox_deltas: (num_rois, num_classes*4)
        targets: list[dict]，每个图像一个标注 dict
        """
        device = cls_logits.device                       # 当前设备
        num_rois = rois.shape[0]                         # 总 RoI 数量
        labels = torch.zeros((num_rois,), dtype=torch.long, device=device)       # 初始化所有 RoI 标签为背景 0
        bbox_targets = torch.zeros((num_rois, 4), dtype=torch.float32, device=device)  # 初始化回归目标为 0

        # 按 图像索引（batch_idx） 分组处理 RoI
        for img_idx, target in enumerate(targets):       # 遍历每张图像的 GT
            gt_boxes = target["boxes"].to(device)        # 取出该图的 GT 框 (Mi,4)
            gt_labels = target["labels"].to(device)      # 取出该图的 GT 类别 (Mi,)

            img_mask = (rois[:, 0] == img_idx)           # 找出属于该图像的 RoI 索引
            props_img = rois[img_mask, 1:5]              # 取出该图像的 proposals (Ni,4)

            # 为该图像的 proposals 分配目标标签和回归目标
            labels_img, bbox_targets_img = assign_targets_to_proposals(
                props_img, gt_boxes, gt_labels, num_classes=self.num_classes
            )

            labels[img_mask] = labels_img                # 填入对应的标签
            bbox_targets[img_mask] = bbox_targets_img    # 填入对应的回归目标

        # 分类损失：所有 RoI 的 CrossEntropyLoss
        cls_loss = self.cls_loss_func(cls_logits, labels)

        # 回归损失：只对正样本（label > 0）计算
        pos_mask = labels > 0                            # 正样本 mask
        if pos_mask.any():                               # 存在正样本时才算回归损失
            labels_pos = labels[pos_mask]                # 取出正样本的类别 (num_pos,)

            # bbox_deltas: (num_rois, num_classes*4) → reshape 为 (num_rois, num_classes, 4)
            bbox_deltas = bbox_deltas.view(-1, self.num_classes, 4)
            # 对每个正样本，选取对应类别的那 4 个回归参数 → (num_pos, 4)
            bbox_deltas_pos = bbox_deltas[pos_mask, labels_pos, :]
            # 取出正样本对应的回归目标 → (num_pos, 4)
            bbox_targets_pos = bbox_targets[pos_mask]

            # 使用 Smooth L1 Loss 计算回归损失
            reg_loss = self.reg_loss_func(bbox_deltas_pos, bbox_targets_pos)
        else:
            # 没有正样本时，回归损失为 0
            reg_loss = torch.tensor(0., device=device)

        # 返回一个包含分类、回归及总损失的字典
        loss_dict = {
            "loss_cls": cls_loss,
            "loss_reg": reg_loss,
            "loss_total": cls_loss + reg_loss
        }
        return loss_dict

    # --------- 推理 & NMS ---------
    def inference(self, rois, cls_logits, bbox_deltas,
                  images_shapes,
                  score_thresh=0.05, nms_thresh=0.5, topk_per_image=100):
        """
        推理模式：根据 RoI 的分类和回归结果，解码并做 NMS，输出最终检测结果
        返回: list[dict]，每个图像一个 dict:
            {
                "boxes":  (K, 4),
                "labels": (K,),
                "scores": (K,)
            }
        """
        device = cls_logits.device                       # 当前设备
        num_rois = rois.shape[0]                         # RoI 数量

        # 对分类 logits 做 softmax，得到每个类别的概率
        cls_probs = F.softmax(cls_logits, dim=1)         # (num_rois, num_classes)

        # proposals 原始框（不含 batch_idx）
        proposals = rois[:, 1:5]                         # (num_rois, 4)

        # bbox_deltas reshape 为 (num_rois, num_classes, 4)
        bbox_deltas = bbox_deltas.view(num_rois, self.num_classes, 4)

        results = []                                     # 存放每张图像的检测结果
        num_images = len(images_shapes)                  # 图像数量

        for img_idx in range(num_images):                # 遍历每一张图像
            h, w = images_shapes[img_idx]                # 取出该图像的高和宽

            img_mask = (rois[:, 0] == img_idx)           # 找出属于该图的 RoI
            if not img_mask.any():                       # 若没有 RoI
                results.append({
                    "boxes": torch.zeros((0, 4), device=device),
                    "labels": torch.zeros((0,), dtype=torch.long, device=device),
                    "scores": torch.zeros((0,), device=device),
                })
                continue                                 # 跳到下一张图

            props_img = proposals[img_mask]              # 该图像的 proposals (Ni,4)
            scores_img = cls_probs[img_mask]             # 该图像的分类概率 (Ni,num_classes)
            deltas_img = bbox_deltas[img_mask]           # 该图像的回归量 (Ni,num_classes,4)

            boxes_all = []                               # 存放该图所有类别的 boxes
            labels_all = []                              # 存放对应的类别标签
            scores_all = []                              # 存放对应的得分

            # 遍历每一个前景类别（1..num_classes-1；0 是背景）
            for cls_id in range(1, self.num_classes):
                scores_cls = scores_img[:, cls_id]       # 取出该类别的得分 (Ni,)
                keep = scores_cls > score_thresh         # 根据 score 阈值做初筛
                if not keep.any():                       # 如果没有得分超过阈值的 RoI
                    continue                             # 跳过这个类别

                scores_cls = scores_cls[keep]            # 筛选后的得分 (M,)
                deltas_cls = deltas_img[keep, cls_id, :] # 对应该类别的回归量 (M,4)
                props_cls = props_img[keep]              # 对应的 proposals (M,4)

                # 解码出预测框坐标
                boxes_cls = decode_boxes(props_cls, deltas_cls)  # (M,4)
                # 将预测框限制在图像边界内
                boxes_cls = clip_boxes_to_image(boxes_cls, h, w)

                # 对该类别的预测框做 NMS
                keep_idx = nms(boxes_cls, scores_cls, nms_thresh)
                boxes_cls = boxes_cls[keep_idx]          # NMS 后的 boxes
                scores_cls = scores_cls[keep_idx]        # NMS 后的 scores
                labels_cls = torch.full((len(keep_idx),),
                                        cls_id,
                                        dtype=torch.long,
                                        device=device)   # 对应类别标签

                boxes_all.append(boxes_cls)              # 保存该类别的 boxes
                labels_all.append(labels_cls)            # 保存该类别的 labels
                scores_all.append(scores_cls)            # 保存该类别的 scores

            if len(boxes_all) == 0:                      # 若所有类别都没有通过阈值
                results.append({
                    "boxes": torch.zeros((0, 4), device=device),
                    "labels": torch.zeros((0,), dtype=torch.long, device=device),
                    "scores": torch.zeros((0,), device=device),
                })
                continue                                 # 继续下一张图

            # 将所有类别的结果拼接起来
            boxes_all = torch.cat(boxes_all, dim=0)      # (K,4)
            labels_all = torch.cat(labels_all, dim=0)    # (K,)
            scores_all = torch.cat(scores_all, dim=0)    # (K,)

            # 按得分从高到低排序
            scores_all, order = scores_all.sort(descending=True)
            boxes_all = boxes_all[order]
            labels_all = labels_all[order]

            # 若超过 topk_per_image，则只保留前 topK
            if boxes_all.shape[0] > topk_per_image:
                boxes_all = boxes_all[:topk_per_image]
                labels_all = labels_all[:topk_per_image]
                scores_all = scores_all[:topk_per_image]

            # 将该图的检测结果打包
            results.append({
                "boxes": boxes_all,
                "labels": labels_all,
                "scores": scores_all
            })

        return results                                   # 返回所有图像的检测结果


# ========= 小测试 =========
if __name__ == "__main__":                               # 仅在直接运行该文件时执行
    # 设备选择：优先 CUDA → 再看是否有 MPS（Apple Silicon）→ 否则用 CPU
    if torch.cuda.is_available():                        # 如果有 NVIDIA GPU（CUDA）
        device = "cuda"
    elif torch.backends.mps.is_available():              # 否则如果有 Apple 的 MPS（Metal）
        device = "mps"
    else:                                                # 否则回退到 CPU
        device = "cpu"

    print("Using device:", device)                       # 打印当前使用的设备

    model = FastRCNN(num_classes=21).to(device)          # 创建模型并移动到指定设备
    model.eval()                                         # 切换到 eval 模式（关闭 dropout、BN 更新）

    # 构造一张假数据：batch 大小为 1，通道为 3，尺寸为 800x600
    images = torch.randn(1, 3, 800, 600, device=device)

    # 构造两个假 RoI：
    # RoI1 属于 batch 中的第 0 张图，框坐标为 (50, 50, 400, 500)
    # RoI2 属于 batch 中的第 0 张图，框坐标为 (100, 120, 500, 700)
    rois = torch.tensor([
        [0, 50.0,  50.0, 400.0, 500.0],
        [0, 100.0, 120.0, 500.0, 700.0],
    ], device=device)

    # 推理模式（不传 targets，会走 inference 流程）
    with torch.no_grad():                               # 关闭梯度计算（推理更省内存）
        detections = model(images, rois, targets=None)  # 得到检测结果

    print(detections)                                   # 打印结果结构（boxes, labels, scores）
